package com.example.harmonet.harmtorch;

public class Tensor {
    private float[] _tensor;
    private int[] _dim;

    public Tensor(int ... dims) {
        int sum = 1;
        for (int dim : dims) {
            sum *= dim;
        }
        _tensor = new float[sum];
        _dim = dims.clone();
    }

    int pos(int ... idx) throws Exception {
        if (idx.length != _dim.length) {
            throw new Exception("index error");
        }
        int pos = 0;
        int size = 1;
        for (int i = idx.length - 1; i >=0; --i) {
            pos += idx[i] * size;
            size *= _dim[i];
        }
        return pos;
    }

    public Tensor view() {
        Tensor rc = new Tensor(_dim);
        rc._tensor = _tensor.clone();
        return rc;
    }

    public Tensor view(int ... dims) {
        Tensor rc = new Tensor(dims);
        rc._tensor = _tensor.clone();
        return rc;
    }

    public void setTensor(float[] vals){
        _tensor = vals.clone();
    }

    public void setTensor(float[][] vals) {
        int dim1 = vals.length;
        int dim2 = vals[0].length;
        _tensor = new float[dim1 * dim2];
        int idx = 0;
        for(int i = 0; i < dim1; ++i) {
            for(int j = 0; j < dim2; ++j) {
                _tensor[idx] = vals[i][j];
                ++idx;
            }
        }
    }

    public void setTensor(float[][][] vals) {
        int dim1 = vals.length;
        int dim2 = vals[0].length;
        int dim3 = vals[0][0].length;
        _tensor = new float[dim1 * dim2 * dim3];
        int idx = 0;
        for(int i = 0; i < dim1; ++i) {
            for(int j = 0; j < dim2; ++j) {
                for(int k = 0; k < dim3; ++k) {
                    _tensor[idx] = vals[i][j][k];
                    ++idx;
                }
            }
        }
    }

    public void setTensorWithNormalization(byte[] vals){
        int tmp = 0;
        for(int i = 0; i < _tensor.length; i++){
            tmp = Byte.toUnsignedInt(vals[i]);
            _tensor[i] = (float)(tmp)/256;
        }
    }

    public void setTensor(float val, int ... coordinates)throws Exception{
        // coordinates: from low dim to high dim
        int pos = 0;
        int coorWeight = 1;
        for(int i = 0; i < coordinates.length; i++){
            if(coordinates[i] >= _dim[i]){
                throw new ArrayIndexOutOfBoundsException("Index of coordinates out of tensor's boundary.");
            }
            else{
                pos += coorWeight * coordinates[i];
                coorWeight *= _dim[i];
            }
        }
        _tensor[pos] = val;
    }

    public String toString() {
        String s = ((Float)_tensor[0]).toString();;
        for(int i = 1; i < _tensor.length; i++){
            s += ", " + ((Float)_tensor[i]).toString();
        }
        return s;
    }

    public float[] tensor(){
        return _tensor;
    }

    public int[] dim(){
        return _dim;
    }
}
